package com.ruoyi.admin.service.document.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.lang.Snowflake;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.ruoyi.admin.domain.LlmDocumentSegments;
import com.ruoyi.admin.domain.LlmDocuments;
import com.ruoyi.admin.mapper.LlmDocumentsSegmentsMapper;
import com.ruoyi.admin.service.document.ILlmDocumentSegmentsService;
import com.ruoyi.common.constant.MilvusConstants;
import com.ruoyi.common.utils.DateUtils;
import com.ruoyi.common.utils.MilvusUtils;
import com.ruoyi.common.utils.StringUtils;
import com.ruoyi.common.utils.embedding.EmbeddingUtils;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 文档切片服务
 */

@Service
@Slf4j
@RequiredArgsConstructor
public class LlmDocumentSegmentsServiceImpl implements ILlmDocumentSegmentsService {


    private final LlmDocumentsSegmentsMapper llmDocumentsSegmentsMapper;

    /**
     * 批量插入
     *
     * @param segmentMeta    元数据
     * @param segmentContent 切片原文
     * @param documents      文档实体
     */
    @Override
    public boolean insert(Map<String, String> segmentMeta,
                          String segmentContent,
                          LlmDocuments documents) {
        //过滤掉值为空的
        segmentMeta = segmentMeta.entrySet().stream().filter(entry -> StringUtils.isNotBlank(entry.getValue()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

        // 创建文段分段对象
        LlmDocumentSegments llmDocumentSegments = new LlmDocumentSegments();
        // 统计字数
        // 拷贝信息
        llmDocumentSegments.setDocumentId(documents.getId());
        llmDocumentSegments.setBusinessData(documents.getBusinessData());
        //新增默认启用
        llmDocumentSegments.setEnabled(0);
        //文档元数据
        String businessData = documents.getBusinessData();
        JSONObject jsonObject = JSONObject.parseObject(businessData);
        //将切片元数据覆盖文档元数据
        jsonObject.putAll(segmentMeta);
        llmDocumentSegments.setBusinessData(jsonObject.toJSONString());
        llmDocumentSegments.setContent(segmentContent);
        llmDocumentSegments.setName(documents.getKnowledgeName());
        llmDocumentSegments.setWordCount(segmentContent.length());
        llmDocumentSegments.setUserId(documents.getUserId());
        llmDocumentSegments.setDatasetCode(documents.getDatasetCode());
        llmDocumentSegments.setKnowledgeId(documents.getKnowledgeId());
        llmDocumentSegments.setCreateTime(DateUtils.dateTimeNow("yyyy-MM-dd HH:mm:ss"));
        llmDocumentSegments.setUpdateTime(DateUtils.dateTimeNow("yyyy-MM-dd HH:mm:ss"));
        String jsonString = jsonObject.toJSONString();
        //当json 字段不为空时候就插入标签
        if (StringUtils.isNotBlank(jsonString)) {
            llmDocumentSegments.setLabelData(jsonString);
        }

        //雪花算法来做id
        Snowflake snowflake = new Snowflake();
        long segmentId = snowflake.nextId();
        llmDocumentSegments.setElementId(segmentId);
        // 操作数据库存入知识库文档分段信息
        llmDocumentsSegmentsMapper.insert(llmDocumentSegments);
        // 操作Milvus向指定知识库中存入数据
        Gson gson = new Gson();
        JsonObject jsonData = new JsonObject();
        //文本向量
        List<Float> embeddings = EmbeddingUtils.createEmbeddings(segmentContent);
        jsonData.addProperty("text", segmentContent);
        jsonData.add("vector", gson.toJsonTree(embeddings));
        //元数据处理
        JsonObject matedata = JsonParser.parseString(jsonObject.toJSONString()).getAsJsonObject();
        matedata.addProperty("time", DateUtils.getTime());
        jsonData.add("metadata", matedata);
        jsonData.addProperty("id", segmentId);
        return MilvusUtils.insertData(MilvusConstants.VECTOR_DOCUMENT_SEGMENT, documents.getDatasetCode(), Collections.singletonList(jsonData));
    }

    @Override
    public List<LlmDocumentSegments> getDocumentsSegments(String id, String keyword) {
        QueryWrapper<LlmDocumentSegments> llmDocumentSegmentsQueryWrapper = new QueryWrapper<>();
        llmDocumentSegmentsQueryWrapper.eq("document_id", id);
        llmDocumentSegmentsQueryWrapper.like(StringUtils.isNotBlank(keyword), "content", keyword);
        return llmDocumentsSegmentsMapper.selectList(llmDocumentSegmentsQueryWrapper);
    }

    @Override
    public boolean edit(LlmDocumentSegments segment) {

        //查询segment 关键数据
        QueryWrapper<LlmDocumentSegments> llmDocumentSegmentsQueryWrapper = new QueryWrapper<>();
        llmDocumentSegmentsQueryWrapper.eq("id", segment.getId());
        llmDocumentSegmentsQueryWrapper.select("element_id", "label_data");
        LlmDocumentSegments llmDocumentSegments = llmDocumentsSegmentsMapper.selectOne(llmDocumentSegmentsQueryWrapper);

        if (llmDocumentSegments == null) {
            return false;
        }

        UpdateWrapper<LlmDocumentSegments> llmDocumentSegmentsUpdateWrapper = new UpdateWrapper<>();
        llmDocumentSegmentsUpdateWrapper.eq("id", segment.getId());
        llmDocumentSegmentsUpdateWrapper.set("content", segment.getContent());
        llmDocumentSegmentsUpdateWrapper.set("update_time", DateUtils.getTime());

        int update = llmDocumentsSegmentsMapper.update(llmDocumentSegmentsUpdateWrapper);
        if (update == 0) {
            return false;
        }
        // 操作Milvus修改文本向量
        Gson gson = new Gson();
        JsonObject jsonData = new JsonObject();
        //文本向量
        List<Float> embeddings = EmbeddingUtils.createEmbeddings(segment.getContent());
        jsonData.addProperty("text", segment.getContent());
        jsonData.add("vector", gson.toJsonTree(embeddings));
        JsonObject matedata = JsonParser.parseString(llmDocumentSegments.getLabelData()).getAsJsonObject();
        matedata.addProperty("time", DateUtils.getTime());
        jsonData.add("metadata", matedata);
        jsonData.addProperty("id", llmDocumentSegments.getElementId());
        if (MilvusUtils.updateData(MilvusConstants.VECTOR_DOCUMENT_SEGMENT, llmDocumentSegments.getDatasetCode(), Collections.singletonList(jsonData))) {
            return true;
        } else {
            throw new RuntimeException("向量修改失败");
        }
    }

    @Override
    public boolean editLabelData(LlmDocumentSegments segment) {
        UpdateWrapper<LlmDocumentSegments> llmDocumentSegmentsUpdateWrapper = new UpdateWrapper<>();
        llmDocumentSegmentsUpdateWrapper.eq("id", segment.getId());
        llmDocumentSegmentsUpdateWrapper.set("label_data", segment.getLabelData());
        int update = llmDocumentsSegmentsMapper.update(llmDocumentSegmentsUpdateWrapper);
        return update > 0;
    }

    @Override
    public List<Long> getIdsLabelData(JSONObject jsonObject) {

        StringBuffer jsonSql = new StringBuffer();
        jsonObject.forEach((key, value) -> {
            String s1 = (String) value;
            jsonSql.append(" and (");
            String[] split = s1.split(",");
            List<String> jsons = new ArrayList<>();
            //构建出查询条件
            for (String s : split) {
                // 如果传入的是空字符的情况就进行额外的处理
                if (" ".equals(s)) {
                    String json = "a.label_data ->>" + " '" + key + "' = ''";
                    jsons.add(json);
                    String json1 = "a.label_data ->>" + " '" + key + "' IS NULL";
                    jsons.add(json1);
                } else if (StringUtils.isNotBlank(s)) {
                    String json = "a.label_data ->>" + " '" + key + "' like '%" + s + "%'";
                    jsons.add(json);
                }
            }
            jsonSql.append(CollectionUtil.join(jsons, " or "));
            jsonSql.append(")");
        });

        String jsonSqlStr = jsonSql.toString().replace("and ()", "");
        List<Long> elementIds = new ArrayList<>();
        if (!jsonObject.isEmpty()) {
            elementIds = llmDocumentsSegmentsMapper.getDocumentSegmentsByLabelData(jsonSqlStr);
            if (elementIds.isEmpty()) {
                return new ArrayList<>();
            }
        }
        return elementIds;
    }

    @Override
    public LlmDocumentSegments detail(String id) {

        if (StringUtils.isBlank(id)) {
            return new LlmDocumentSegments();
        }
        return llmDocumentsSegmentsMapper.selectById(id);
    }


}
